import math

from data_reader import DataReader
from heapq import *

class KNN():
  def __init__(self, trans_trial_file, trans_train_file, N=3):
    test_reader = DataReader(trans_trial_file)
    train_reader = DataReader(trans_train_file)
    
    test_feature_vectors = test_reader.get_feature_vectors()
    train_feature_vectors = train_reader.get_feature_vectors()

    self.neighbours = []

    # Collect the distances in a heap and select the top 3
    for test_feature_vector in test_feature_vectors:
      heap = []
      train_idx = 0
      for train_feature_vector in train_feature_vectors:
        heappush(heap,
        (self.distance(test_feature_vector, train_feature_vector),
         train_feature_vector[1], train_idx))
        train_idx += 1
      cur_neighbours = [entry[2] for entry in nsmallest(N, heap)]
      self.neighbours.append(cur_neighbours)

  # Compute Euclidian distance between vector a and b
  def distance(self, vector_a, vector_b):
    cur_distance = 0
    for i in range(0, len(vector_a[0])):
      cur_distance += (int(vector_a[0][i]) - int(vector_b[0][i])) ** 2
    return math.sqrt(cur_distance)

  # Returns the neighbours for the test file.
  def get_neighbours(self):
    return self.neighbours
